import math
import numpy as np

import paddle.fluid.layers as L
from paddle.fluid.dygraph import Layer, Sequential, Conv2D
from helm.static.models.layers import Conv2d, BN
from helm.dynamic.models.init import init_zeros_, init_normal_
from helm.static.models.modules import SELayer


# Residual Channel Attention Block (RCAB)
class RCAB(Layer):

    def __init__(self, in_channels, reduction=16, bn=False):
        super(RCAB, self).__init__()
        modules_body = [
            Conv2d(in_channels, in_channels, 3, bn=bn, act='default'),
            Conv2d(in_channels, in_channels, 3, bn=bn),
            SELayer(in_channels, reduction)
        ]
        if bn:
            init_zeros_(modules_body[1][1].weight)
        self.body = Sequential(*modules_body)

    def forward(self, x):
        return x + self.body(x)


# Residual Group (RG)
class ResidualGroup(Layer):

    def __init__(self, in_channels, reduction, n_resblocks, bn=False):
        super(ResidualGroup, self).__init__()
        modules_body = [
            RCAB(in_channels, reduction, bn=bn)
            for _ in range(n_resblocks)
        ]
        modules_body.append(Conv2d(in_channels, in_channels, 3, bn=bn))
        if bn:
            init_zeros_(modules_body[-1][1].weight)
        self.body = Sequential(*modules_body)

    def forward(self, x):
        return x + self.body(x)


class Scale(Layer):
    def __init__(self, mean, std):
        super().__init__()
        mean = np.array(mean)
        std = np.array(std)
        assert mean.shape == std.shape
        assert len(mean.shape) == 1
        c = mean.shape[0]
        mean = np.array(mean).reshape([1, c, 1, 1]).astype(np.float32)
        std = np.array(std).reshape([1, c, 1, 1]).astype(np.float32)
        self.mean = self.create_parameter([1, c, 1, 1])
        self.mean.set_value(mean)
        self.mean.trainable = False
        self.std = self.create_parameter([1, c, 1, 1])
        self.std.set_value(std)
        self.std.trainable = False

    def forward(self, x):
        return (x + self.mean) * self.std


class Unscale(Layer):
    def __init__(self, mean, std):
        super().__init__()
        mean = np.array(mean)
        std = np.array(std)
        c = mean.shape[0]
        mean = np.array(mean).reshape([1, c, 1, 1]).astype(np.float32)
        std = np.array(std).reshape([1, c, 1, 1]).astype(np.float32)
        self.mean = self.create_parameter([1, c, 1, 1])
        self.mean.set_value(mean)
        self.mean.trainable = False
        self.std = self.create_parameter([1, c, 1, 1])
        self.std.set_value(std)
        self.std.trainable = False

    def forward(self, x):
        return x / self.std - self.mean


class PixelShuffle(Layer):

    def __init__(self, upscale_factor):
        super().__init__()
        self.upscale_factor = upscale_factor

    def forward(self, x):
        return L.pixel_shuffle(x, self.upscale_factor)


class Upsampler(Sequential):

    def __init__(self, scale, channels, bn=False):
        super().__init__()
        m = []
        if (scale & (scale - 1)) == 0:  # Is scale = 2^n?
            for _ in range(int(math.log(scale, 2))):
                m.append(Conv2d(channels, channels * 4, 3))
                m.append(PixelShuffle(2))
                if bn:
                    m.append(BN(channels))
        elif scale == 3:
            m.append(Conv2d(channels, 9 * channels, 3))
            m.append(PixelShuffle(3))
            if bn:
                if bn: m.append(BN(channels))
        else:
            raise NotImplementedError
        super().__init__(*m)


class RCAN(Layer):

    def __init__(self, mean, std, n_resgroups, n_resblocks, channels, scale, bn=False, init_std=0.02):
        super().__init__()
        reduction = 16

        # DIV2K 1-800 (0.4488, 0.4371, 0.4040)
        # DIVFlickr2K 1-3450 (0.4690, 0.4490, 0.4036)

        self.sub_mean = Unscale(mean, std)

        modules_head = [Conv2d(3, channels, 3, bn=bn)]
        modules_body = [
            ResidualGroup(channels, reduction, n_resblocks, bn=bn)
            for _ in range(n_resgroups)
        ]

        modules_body.append(Conv2d(channels, channels, 3, bn=bn))
        if bn:
            init_zeros_(modules_body[-1][1].weight)

        modules_tail = [
            Upsampler(scale, channels, bn=bn),
            Conv2d(channels, 3, 3)
        ]

        self.add_mean = Scale(mean, std)

        self.head = Sequential(*modules_head)
        self.body = Sequential(*modules_body)
        self.tail = Sequential(*modules_tail)

        for m in self.sublayers():
            if isinstance(m, Conv2D):
                init_normal_(m.weight, 0, init_std)
                init_zeros_(m.bias)

    def forward(self, x):
        x = self.sub_mean(x)
        x = self.head(x)

        res = self.body(x)
        res += x

        x = self.tail(res)
        x = self.add_mean(x)
        return x
